import numpy as np
import pandas as pd
import torch
from scipy.optimize import minimize
import os
import time

'''
0.作者：张哲

1.功能介绍
    这是一个用不同的最紧密堆积计算出的理论曲线与实际曲线的拟合优度的计算包。

2.参数介绍
    2.1 Dinger_Funk方法：调用Dinger_Funk最紧密堆积理论
    


'''

os.environ['CUDA_LAUNCH_BLOCKING'] = '1'

def check_cuda():
    if torch.cuda.is_available():
        print("CUDA Available:", torch.cuda.is_available())
        print("CUDA Device Name:", torch.cuda.get_device_name(0))
        print("CUDA Device Count:", torch.cuda.device_count())
    else:
        print("CUDA is not available. Using CPU.")

def compute_ideal_distribution(D, n=0.23, device='cpu'):
    D_max = max(D)
    D_min = min(D)
    ideal_distribution = [((d ** n - D_min ** n) / (D_max ** n - D_min ** n)) * 100 for d in D]
    if device == 'cuda':
        return torch.tensor(ideal_distribution, device='cuda')
    else:
        return torch.tensor(ideal_distribution, device='cpu')

def objective_function(v, data_tensors, ideal_distribution):
    act = sum(v[i] * data_tensors[i] for i in range(len(v)))
    r2 = torch.corrcoef(torch.stack([ideal_distribution, act]))[0, 1].item()
    return -r2  #

def find_best_combination_optimized(combinations_df, row_data, ideal_distribution, num_samples, max_bound=1, device='cpu'):
    best_combination_data = []
    for index, row in combinations_df.iterrows():
        samples = [row[f'Sample {i + 1}'] for i in range(num_samples)]
        if device == 'cuda':
            data_tensors = [torch.tensor(row_data.loc[sample].values, device='cuda') for sample in samples]
        else:
            data_tensors = [torch.tensor(row_data.loc[sample].values, device='cpu') for sample in samples]

        initial_guess = np.ones(num_samples) / num_samples
        constraints = [{'type': 'eq', 'fun': lambda v: np.sum(v) - 1}]
        bounds = [(0.01, max_bound) for _ in range(num_samples)]
        
        result = minimize(objective_function, initial_guess, args=(data_tensors, ideal_distribution),
                          constraints=constraints, bounds=bounds, method='SLSQP')
        V = [round(result.x[i], 4) for i in range(num_samples)]
        best_combination_data.append({
            'Combination': tuple(samples),
            'R2': -result.fun,
            "V": V,
            'Distribution': [round(float(x), 2) for x in (sum(result.x[i] * data_tensors[i] for i in range(num_samples))).cpu().numpy()]
        })
    return best_combination_data



def Dinger_Funk(file_name, n_comps, n=0.23, device=0):
    '''
    1.Dinger Funk最紧密堆积理论

    2.Parameters:

        n_comps: 需要计算几元复配
            例如：当n_comps=3时，计算的是3元复配方案的结果。需要注意的是在运行本脚本之前，必须先运行Sample_Combinations脚本，以生成不同元的所有组合方案

        n: 分布模数。
            一般为0-1之间的小数，默认为0.23.

        device: 用来计算的设备.
            0：使用CPU计算；
            1：使用GPU计算。

    3.return:
            生成一个名为Optimized Method Results的文件夹，文件夹内保存了指定元计算的结果。

    '''
    print(f"您当前计算的是{n_comps}元复配")
    
    if device == 1:
        check_cuda()
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    else:
        device = 'cpu'

    row_data = pd.read_excel(file_name, index_col=0, header=0)
    D = list(map(float, row_data.columns))
    print(f"D分布为: {D}")
    ideal_distribution = compute_ideal_distribution(D, n=n, device=device)

    combinations_file = f'sample_combinations/{n_comps}_sample_combinations.csv'
    combinations_df = pd.read_csv(combinations_file)

    start_time = time.time()
    best_combination_data_optimized = find_best_combination_optimized(combinations_df, row_data, ideal_distribution, num_samples=n_comps, device=device)
    optimized_time = time.time() - start_time
    print(f"Optimized method time: {optimized_time:.2f} seconds")

    # Ensure results folder exists
    results_folder = 'Optimized Method Results'
    if not os.path.exists(results_folder):
        os.makedirs(results_folder)

    output_file = os.path.join(results_folder, f'best_combinations_{n_comps}_optimized.csv')
    df_best_combinations_optimized = pd.DataFrame(best_combination_data_optimized)
    df_best_combinations_optimized.to_csv(output_file, index=False, float_format='%.4f')

    print("Optimized Method Results:")
    print(df_best_combinations_optimized)


